package com.lj.common.utils;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.thread.GlobalThreadPool;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * @author luojing
 * @since 2025/3/19 14:13
 * 批处理工具
 */
public class BatchDisposeUtils {

    /**
     * 批量处理任务，会等待所有任务处理完毕
     *
     * @param taskCollection 任务集合
     * @param execFun        任务执行函数
     * @param <T>            任务的类型，也是任务执行函数的入参
     */
    public static <T> void dispose(Collection<T> taskCollection, Consumer<T> execFun) throws InterruptedException {
        dispose(taskCollection, execFun, GlobalThreadPool.getExecutor(), false);
    }

    /**
     * 批量处理任务，会等待所有任务处理完毕
     *
     * @param taskCollection 任务集合
     * @param execFun        任务执行函数
     * @param executor       任务执行的线程池
     * @param shutdown       执行结束是否关闭线程池
     * @param <T>            任务的类型，也是任务执行函数的入参
     */
    public static <T> void dispose(Collection<T> taskCollection, Consumer<T> execFun, ExecutorService executor, boolean shutdown) throws InterruptedException {
        dispose(taskCollection, t -> {
            execFun.accept(t);
            return null;
        }, executor, shutdown);
    }

    /**
     * 批量处理任务，会等待所有任务处理完毕获取结果
     *
     * @param taskCollection 任务集合
     * @param execFun        任务执行函数
     * @param <T>            任务的类型，也是任务执行函数的入参
     * @param <R>            任务的结果
     * @return 执行结果的Future对象
     */
    public static <T, R> List<Future<R>> dispose(Collection<T> taskCollection, Function<T, R> execFun) throws InterruptedException {
        return dispose(taskCollection, execFun, GlobalThreadPool.getExecutor(), false);
    }

    /**
     * 批量处理任务，会等待所有任务处理完毕获取结果
     *
     * @param taskCollection 任务集合
     * @param execFun        任务执行函数
     * @param executor       任务执行的线程池
     * @param shutdown       执行结束是否关闭线程池
     * @param <T>            任务的类型，也是任务执行函数的入参
     * @param <R>            任务的结果
     * @return 执行结果的Future对象
     */
    public static <T, R> List<Future<R>> dispose(Collection<T> taskCollection, Function<T, R> execFun, ExecutorService executor, boolean shutdown) throws InterruptedException {
        if (CollUtil.isEmpty(taskCollection)) {
            return Collections.emptyList();
        }
        // 使用CountDownLatch来阻塞当前线程，等待所有任务执行完成后继续执行
        CountDownLatch countDownLatch = new CountDownLatch(taskCollection.size());
        List<Future<R>> result = new ArrayList<>();
        for (T task : taskCollection) {
            Future<R> future = executor.submit(() -> {
                try {
                    return execFun.apply(task);
                } finally {
                    countDownLatch.countDown();
                }
            });
            result.add(future);
        }
        countDownLatch.await();
        if (shutdown) {
            executor.shutdown();
        }
        return result;
    }
}
